import os
from abc import ABC, abstractmethod
from typing import List, Optional, Dict, Any
from dataclasses import dataclass
from pydantic import BaseModel, Field
import httpx
from loguru import logger


@dataclass
class RerankDocument:
    """重排序文档实体"""
    text: str
    index: int
    relevance_score: float
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        return {
            "document": {"text": self.text},
            "index": self.index,
            "relevance_score": self.relevance_score
        }


class RerankResult(BaseModel):
    """重排序结果统一实体"""
    documents: List[RerankDocument] = Field(default_factory=list)
    total_count: int = 0
    model_name: Optional[str] = None
    
    def get_texts(self) -> List[str]:
        """获取排序后的文本列表"""
        return [doc.text for doc in self.documents]
    
    def get_top_n(self, n: int) -> List[RerankDocument]:
        """获取前N个结果"""
        return self.documents[:n]
    
    def filter_by_score(self, min_score: float) -> List[RerankDocument]:
        """根据相关性分数过滤结果"""
        return [doc for doc in self.documents if doc.relevance_score >= min_score]


class RerankLLM(ABC):
    """重排序模型抽象基类"""
    
    @abstractmethod
    def rerank(self, query: str, documents: List[str], top_n: int = 3, 
               min_relevance_score: float = 0.0) -> RerankResult:
        """
        重排序文档
        
        Args:
            query: 查询文本
            documents: 待排序的文档列表
            top_n: 返回前N个结果
            min_relevance_score: 最小相关性分数阈值
            
        Returns:
            RerankResult: 重排序结果
        """
        pass


class BailianRerankLLM(RerankLLM):
    """百炼重排序模型实现"""
    
    def __init__(self, 
                 api_key: Optional[str] = None, 
                 model: str = "gte-rerank-v2", 
                 base_url: str = "https://dashscope.aliyuncs.com/api/v1/services/rerank/text-rerank/text-rerank",
                 timeout: int = 30):
        """
        初始化百炼重排序模型
        
        Args:
            api_key: API密钥，如果为None则从环境变量DASHSCOPE_API_KEY获取
            model: 模型名称，默认为gte-rerank-v2
            base_url: API基础URL
            timeout: 请求超时时间（秒）
        """
        self.api_key = api_key or os.getenv("DASHSCOPE_API_KEY")
        if not self.api_key:
            raise ValueError("API key is required. Please provide api_key parameter or set DASHSCOPE_API_KEY environment variable.")
        
        self.model = model
        self.base_url = base_url
        self.timeout = timeout
        
    def rerank(self, query: str, documents: List[str], top_n: int = 3, 
               min_relevance_score: float = 0.0) -> RerankResult:
        """
        使用百炼API进行文档重排序
        
        Args:
            query: 查询文本
            documents: 待排序的文档列表
            top_n: 返回前N个结果
            min_relevance_score: 最小相关性分数阈值
            
        Returns:
            RerankResult: 重排序结果
        """
        if not documents:
            return RerankResult(total_count=0, model_name=self.model)
        
        # 构建请求数据
        request_data = {
            "model": self.model,
            "input": {
                "query": query,
                "documents": documents
            },
            "parameters": {
                "return_documents": True,
                "top_n": min(top_n, len(documents))
            }
        }
        
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        try:
            with httpx.Client(timeout=self.timeout) as client:
                response = client.post(
                    self.base_url,
                    headers=headers,
                    json=request_data
                )
                response.raise_for_status()
                
                result_data = response.json()
                return self._parse_response(result_data, min_relevance_score)
                
        except httpx.HTTPError as e:
            logger.error(f"HTTP error during rerank request: {e}")
            raise
        except Exception as e:
            logger.error(f"Error during rerank request: {e}")
            raise
    
    def _parse_response(self, response_data: Dict[str, Any], 
                       min_relevance_score: float = 0.0) -> RerankResult:
        """
        解析API响应数据
        
        Args:
            response_data: API响应数据
            min_relevance_score: 最小相关性分数阈值
            
        Returns:
            RerankResult: 解析后的重排序结果
        """
        try:
            # 检查响应状态
            if response_data.get("code") and response_data["code"] != "200":
                error_msg = response_data.get("message", "Unknown error")
                raise ValueError(f"API returned error: {error_msg}")
            
            # 解析输出结果
            output = response_data.get("output", {})
            results = output.get("results", [])
            
            # 构建文档列表
            rerank_documents = []
            for item in results:
                relevance_score = item.get("relevance_score", 0.0)
                
                # 应用最小相关性分数过滤
                if relevance_score >= min_relevance_score:
                    document_data = item.get("document", {})
                    text = document_data.get("text", "")
                    index = item.get("index", 0)
                    
                    rerank_documents.append(RerankDocument(
                        text=text,
                        index=index,
                        relevance_score=relevance_score
                    ))
            
            # 按相关性分数降序排序
            rerank_documents.sort(key=lambda x: x.relevance_score, reverse=True)
            
            return RerankResult(
                documents=rerank_documents,
                total_count=len(rerank_documents),
                model_name=self.model
            )
            
        except Exception as e:
            logger.error(f"Error parsing rerank response: {e}")
            raise ValueError(f"Failed to parse rerank response: {e}")
